package com.gallop.mq.core.template;

import com.gallop.mq.MQProperties;
import com.gallop.mq.MQTemplate;
import com.gallop.mq.core.MQExecutorProvider;
import com.gallop.mq.core.MQMessageCodec;
import com.gallop.mq.core.MQMessageDispatcher;
import com.gallop.mq.core.dispatcher.DispatchContext;
import com.gallop.mq.core.exception.MQException;
import com.gallop.mq.core.template.handler.DeliveryContext;
import com.gallop.mq.core.template.handler.DeliveryType;
import com.gallop.mq.core.template.handler.MessageDeliveryHandler;
import com.gallop.mq.core.template.handler.registry.DynamicProxyListenerCreator;
import com.gallop.mq.core.template.handler.registry.MessageDeliveryHandlerRegistry;
import com.gallop.mq.core.template.handler.registry.RegistryHolder;
import com.gallop.mq.core.template.handler.registry.RegistryPreprocessor;
import com.google.common.collect.Maps;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.thread.TaskRunnerFactory;
import org.springframework.cglib.proxy.Enhancer;
import org.springframework.cglib.proxy.MethodInterceptor;
import org.springframework.cglib.proxy.MethodProxy;

import javax.jms.*;
import java.io.ByteArrayOutputStream;
import java.lang.reflect.Method;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

@Getter
@Setter
@Slf4j
public class ActiveMQTemplate implements MQTemplate, RegistryHolder {

    private MQMessageCodec codec;
    private MQMessageDispatcher dispatcher;
    private MQProperties properties;
    private MQExecutorProvider executorProvider;
    private TaskRunnerFactoryProxy taskRunnerFactoryProxy = new TaskRunnerFactoryProxy();
    private Connection connection;
    private Session session;

    private Queue localReplyQueue;
    private String localReplyQueueName;

    private Map<String, PendingReply> pendingReplyMap = Maps.newConcurrentMap();

    private MessageDeliveryHandlerRegistry handlerRegistry;

    @Override
    public void init(MQProperties properties, MQMessageDispatcher dispatcher, MQMessageCodec codec, MQExecutorProvider executorProvider, MessageDeliveryHandlerRegistry handlerRegistry) throws MQException {
        setProperties(properties);
        setDispatcher(dispatcher);
        setCodec(codec);
        setExecutorProvider(executorProvider);

        this.initRegistry(handlerRegistry);

        ActiveMQConnectionFactory connectionFactory = new ActiveMQConnectionFactory();
        String failoverURL = String.format("failover:(%s)?randomize=false", properties.getHosts().get(0));
        connectionFactory.setBrokerURL(failoverURL);
        connectionFactory.setUserName(properties.getUser());
        connectionFactory.setPassword(properties.getSecret());
        connectionFactory.setSessionTaskRunner(taskRunnerFactoryProxy.enhance(new TaskRunnerFactory()));

        wrapJMSException(() -> {
            this.connection = connectionFactory.createConnection();
            this.session = this.connection.createSession(false, Session.AUTO_ACKNOWLEDGE);

            // 初始化rpc接收队列
            this.localReplyQueueName = "local-reply-" + UUID.randomUUID().toString().replaceAll("-", "");
            localReplyQueue = session.createQueue(localReplyQueueName);
            MessageConsumer localReplyConsumer = this.session.createConsumer(localReplyQueue);
            localReplyConsumer.setMessageListener(handlerRegistry.createListener(DeliveryContext.builder().mqTemplate(this).build()));
            // 初始化扫描监听
            dispatcher.getAllContainers().forEach(container -> wrapJMSException(() -> {
                Topic topic = this.session.createTopic(container.getTopic());
                MessageConsumer topicConsumer = this.session.createConsumer(topic);
                topicConsumer.setMessageListener(handlerRegistry.createListener(DeliveryContext.builder()
                        .mqTemplate(this).container(container).build()));
            }, "ActiveMQ consumer create failed."));
            this.connection.start();
        }, "ActiveMQ init failed.");
    }

    private Object decodeMessage(Message message) {
        BytesMessage transferMessage = (BytesMessage) message;
        int bufferSize = properties.getAmq().getIoBufferSize();
        try (ByteArrayOutputStream out = new ByteArrayOutputStream(bufferSize)) {
            byte[] buffer = new byte[bufferSize];
            int bytesRead;
            while ((bytesRead = transferMessage.readBytes(buffer)) > 0) {
                out.write(buffer, 0, bytesRead);
            }
            out.flush();
            byte[] bytes = out.toByteArray();
            return this.codec.decode(bytes);
        } catch (Exception e) {
            throw new MQException("parse ActiveMQ Message error.", e);
        }
    }

    @Override
    public void onDispose() {
        log.info("ActiveMQ destroy");
        try {
            this.session.close();
            this.connection.close();
        } catch (JMSException e) {
            log.error("", e);
        }
    }

    @Override
    public void send(Object message, String topic) {
        try {
            Queue queue = session.createQueue(topic);
            MessageProducer producer = this.session.createProducer(queue);
            BytesMessage transferMessage = this.session.createBytesMessage();
            byte[] bytes = this.codec.encode(message);
            transferMessage.setJMSType(DeliveryType.SEND.name());
            transferMessage.writeBytes(bytes);
            producer.send(transferMessage);
        } catch (Exception ex) {
            throw new MQException("MQ(ActiveMQ) send failed.", ex);
        }
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse) {
        request(message, topic, requestTimeout, onResponse, null);
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse, BiConsumer<Object, Exception> onError) {
        PendingReply pendingReply = createPendingReply();
        try {
            P ret = doRequest(message, topic, pendingReply, requestTimeout);
            onResponse.accept(ret);
        } catch (Exception ex) {
            log.error("MQ(ActiveMQ) send failed.", ex);
            pendingReplyMap.remove(pendingReply.getSavedCorrelation());
            pendingReply.returned(ex);
            if (onError != null) {
                onError.accept(message, ex);
            }
        }
    }

    @Override
    public <P> P request(Object message, String topic, Duration receiveTimeout) throws InterruptedException {
        PendingReply pendingReply = createPendingReply();
        try {
            return doRequest(message, topic, pendingReply, receiveTimeout);
        } catch (Exception ex) {
            log.error("MQ(ActiveMQ) request failed.", ex);
            pendingReplyMap.remove(pendingReply.getSavedCorrelation());
            pendingReply.returned(ex);
            return null;
        }
    }

    private PendingReply createPendingReply() {
        String savedCorrelation = "rpc-correlation-" + UUID.randomUUID().toString().replaceAll("-", "");
        PendingReply pendingReply = new PendingReply();
        pendingReply.setSavedCorrelation(savedCorrelation);
        pendingReply.setSavedReplyTo(localReplyQueueName);
        pendingReplyMap.put(savedCorrelation, pendingReply);
        return pendingReply;
    }

    private <P> P doRequest(Object message, String topic, PendingReply pendingReply, Duration receiveTimeout) throws Exception {
        String savedCorrelation = pendingReply.getSavedCorrelation();

        BytesMessage transferMessage = this.session.createBytesMessage();
        byte[] reqBytes = this.codec.encode(message);
        transferMessage.writeBytes(reqBytes);
        transferMessage.setJMSReplyTo(localReplyQueue);
        transferMessage.setJMSCorrelationID(savedCorrelation);
        transferMessage.setJMSType(DeliveryType.REQUEST.name());

        Topic sendTopic = session.createTopic(topic);
        MessageProducer producer = this.session.createProducer(sendTopic);
        producer.send(transferMessage);

        if (receiveTimeout == null && properties.getRequestTimeout() > 0) {
            receiveTimeout = Duration.ofMillis(properties.getRequestTimeout());
        }
        if (receiveTimeout == null) {
            return pendingReply.get();
        }
        return pendingReply.get(receiveTimeout.toMillis(), TimeUnit.MILLISECONDS);
    }

    @Override
    public MessageDeliveryHandlerRegistry.InitializerBuilder registryInitializer() {
        MessageDeliveryHandlerRegistry.InitializerBuilder initializerBuilder = MessageDeliveryHandlerRegistry.InitializerBuilder.newBuilder();
        initPreprocessor(initializerBuilder);
        initHandler(initializerBuilder);
        return initializerBuilder;
    }

    private void initPreprocessor(MessageDeliveryHandlerRegistry.InitializerBuilder initializerBuilder) {
        initializerBuilder.setPreprocessor(RegistryPreprocessor.<MessageListener, Message>builder()
                // 使用动态代理调用转发逻辑，此处可自定义functionInterface，在回调中手动触发handlerRegistry#handleDelivery
                .listenerCreator(DynamicProxyListenerCreator.ofInterface(handlerRegistry, MessageListener.class))
                .deliveryTypeProvider(args -> {
                    AtomicReference<DeliveryType> deliveryType = new AtomicReference<>();
                    wrapJMSException(() -> {
                        Message message = args.getArgAs(0);
                        deliveryType.set(DeliveryType.valueOf(message.getJMSType()));
                    });
                    return deliveryType.get();
                })
                .rpcInfoProvider(args -> {
                    MessageDeliveryHandler.CustomRpcInfo.CustomRpcInfoBuilder builder = MessageDeliveryHandler.CustomRpcInfo.builder();
                    wrapJMSException(() -> {
                        Message message = args.getArgAs(0);
                        builder.correlationId(message.getJMSCorrelationID());
                        // custom replyTo
                        builder.replyTo(message.getJMSReplyTo());
                    });
                    return builder.build();
                })
                .rawMessageProvider(args -> args.getArgAs(0))
                .mqMessageProvider(this::decodeMessage)
                .build());
    }

    private void initHandler(MessageDeliveryHandlerRegistry.InitializerBuilder initializerBuilder) {
        initializerBuilder.registerHandler(DeliveryType.RESPONSE, (deliveryContext, deliveryType, rpcInfo, msg) -> wrapJMSException(() -> {
                    // 处理RPC响应
                    PendingReply reply = pendingReplyMap.remove(rpcInfo.getCorrelationId());
                    if (reply == null) {
                        log.warn("rpc PendingReply not found, response will be ignored. correlationId={}, replyTo={},response={}",
                                rpcInfo.getCorrelationId(), rpcInfo.getReplyTo(), msg);
                        return;
                    }
                    // future complete
                    reply.reply(msg);
                }, "onResponse failed."))
                .registerHandler(DeliveryType.SEND, (deliveryContext, deliveryType, rpcInfo, msg) -> wrapJMSException(() -> {
                    DispatchContext.DispatchContextBuilder contextBuilder = DispatchContext.builder()
                            .topic(deliveryContext.getContainer().getTopic())
                            .group(deliveryContext.getContainer().getGroup())
                            .message(msg);
                    this.dispatcher.dispatch(contextBuilder.build());
                }, "ActiveMQ consumer handler failed."))
                .registerHandler(DeliveryType.REQUEST, (deliveryContext, deliveryType, rpcInfo, msg) -> wrapJMSException(() -> {
                    DispatchContext.DispatchContextBuilder contextBuilder = DispatchContext.builder()
                            .topic(deliveryContext.getContainer().getTopic())
                            .group(deliveryContext.getContainer().getGroup())
                            .message(msg);
                    // rpc response
                    contextBuilder.replyTo(rpcInfo.getCorrelationId()).correlationId(rpcInfo.getCorrelationId())
                            .rpcResponse((correlationId, replyTo, response) -> wrapJMSException(() -> {
                                Destination destination = (Destination) rpcInfo.getReplyTo();
                                BytesMessage responseMessage = this.session.createBytesMessage();
                                responseMessage.writeBytes(this.codec.encode(response));
                                responseMessage.setJMSCorrelationID(correlationId);
                                responseMessage.setJMSReplyTo(destination);
                                responseMessage.setJMSType(DeliveryType.RESPONSE.name());
                                MessageProducer producer = this.session.createProducer(destination);
                                producer.send(responseMessage);
                            }, "MQ(ActiveMQ) response error."));
                    this.dispatcher.dispatch(contextBuilder.build());
                }, "ActiveMQ consumer handler failed."));
    }

    @FunctionalInterface
    private interface JMSThrowableLogic {
        void run() throws JMSException;
    }

    private static void wrapJMSException(JMSThrowableLogic code) {
        wrapJMSException(code, "");
    }

    private static void wrapJMSException(JMSThrowableLogic code, String errorMessage) {
        try {
            code.run();
        } catch (JMSException exception) {
            throw new MQException(errorMessage, exception);
        }
    }

    /**
     * 替换TaskRunnerFactory的线程池
     */
    private class TaskRunnerFactoryProxy {

        private TaskRunnerFactoryProxy() {
        }

        public TaskRunnerFactory enhance(TaskRunnerFactory src) {
            Enhancer enhancer = new Enhancer();
            enhancer.setSuperclass(src.getClass());
            enhancer.setCallback(new TaskRunnerFactoryInterceptor(src));
            return (TaskRunnerFactory) enhancer.create();
        }

        private class TaskRunnerFactoryInterceptor implements MethodInterceptor {

            private final TaskRunnerFactory sourceObject;

            public TaskRunnerFactoryInterceptor(TaskRunnerFactory sourceObject) {
                this.sourceObject = sourceObject;
            }

            @Override
            public Object intercept(Object o, Method method, Object[] objects, MethodProxy methodProxy) throws Throwable {
                if (method.getName().equals("createDefaultExecutor")) {
                    return executorProvider.create();
                }
                method.setAccessible(true);
                return method.invoke(sourceObject, objects);
            }
        }
    }
}
